'''
decode raw data
'''

import numpy as np
import matplotlib.pyplot as plt

root = "/home/nuc/workspace/ai-robot-plus/catkin_ws/src/xrobot/data/"

# list all files and sort by name
import os
files = os.listdir(root)
# find all txt files
files = [f for f in files if f.endswith(".txt")]
files.sort()
# use the last file
file = root + files[-1]

# file = "/home/nuc/workspace/ai-robot-plus/catkin_ws/src/xrobot/data/data_20230507_162317.txt"
print("file: ", file)
data = np.loadtxt(file, dtype=np.float64)
print(data.shape)

# 1~3: gyro
# 4~6: acc
# 7~10: encoder increment

time = data[:, 0] * 0.001 # ms -> s
gyro_rps = data[:, 1:4] * 0.0010650782 # ±2000dps (4000 * 3.1415926535 / 180 / 65536)
acc_mps2 = data[:, 4:7] * 0.004790 # ±16g (32 * 9.81 / 65536) 
enc_rpm = data[:, 7:11] / 30000.0 / 0.001 * 60.0 # 30000 pulse per round, 0.001s per loop, rpm

# macanum wheel decode
wheel_circum = 0.2356194 # diameter 75mm
vel_x = ( enc_rpm[:, 0] + enc_rpm[:, 1] + enc_rpm[:, 2] + enc_rpm[:, 3]) / 4.0 / 60.0 * wheel_circum # m/s
vel_y = (-enc_rpm[:, 0] + enc_rpm[:, 1] + enc_rpm[:, 2] - enc_rpm[:, 3]) / 4.0 / 60.0 * wheel_circum # m/s
# need wheels distance
# omega_z = (-enc_rpm[:, 0] + enc_rpm[:, 1] - enc_rpm[:, 2] + enc_rpm[:, 3]) / 4.0 / 60.0 * wheel_circum # rad/s

# combine all data
data = np.hstack((time.reshape(-1, 1), gyro_rps, acc_mps2, vel_x.reshape(-1, 1), vel_y.reshape(-1, 1)))
# save to file

# split filename by "."
result_file = root + "decode/" + files[-1].split(".")[0] + "_decode.txt"

np.savetxt(result_file, data, delimiter=" ", fmt="%.6f")


plt.figure(0)
plt.plot(time, gyro_rps[:, 0])
plt.plot(time, gyro_rps[:, 1])
plt.plot(time, gyro_rps[:, 2])
plt.xlabel("time(s)")
plt.ylabel("gyro(rad/s)")
plt.title('gyro_rps')

plt.figure(1)
plt.plot(time, acc_mps2[:, 0])
plt.plot(time, acc_mps2[:, 1])
plt.plot(time, acc_mps2[:, 2])
plt.xlabel("time(s)")
plt.ylabel("accel(m/s2)")
plt.title('acc_mps2')

plt.figure(2)
plt.plot(time, vel_x)
plt.plot(time, vel_y)
plt.title('vel')
plt.xlabel("time(s)")
plt.ylabel("vel(m/s)")
plt.show()
